import numpy as np
import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import proj3d
from matplotlib.patches import FancyArrowPatch

# ---------------- 全局配置 ----------------
plt.rcParams["font.family"] = ["SimHei", "WenQuanYi Micro Hei", "Heiti TC", "Times New Roman"]
plt.rcParams["axes.unicode_minus"] = False
plt.rcParams.update({
    "font.size": 12,
    "axes.labelsize": 14,
    "axes.titlesize": 16,
    "legend.fontsize": 11,
    "grid.linestyle": "--",
    "grid.alpha": 0.7
})

# LaTeX 设置（容错）
try:
    plt.rcParams["text.usetex"] = True
    plt.rcParams["text.latex.preamble"] = [
        r"\usepackage{amsmath}",
        r"\usepackage{amsfonts}",
        r"\usepackage{amsbsy}"
    ]
except Exception:
    plt.rcParams["text.usetex"] = False


# ---------------- 自定义3D箭头 ----------------
class Arrow3D(FancyArrowPatch):
    def __init__(self, xs, ys, zs, *args, **kwargs):
        super().__init__((0, 0), (0, 0), *args, **kwargs)
        self._verts3d = xs, ys, zs

    def do_3d_projection(self, renderer=None):
        xs3d, ys3d, zs3d = self._verts3d
        xs, ys, zs = proj3d.proj_transform(xs3d, ys3d, zs3d, self.axes.M)
        self.set_positions((xs[0], ys[0]), (xs[1], ys[1]))
        return np.min(zs)


def plot_missile_m1_trajectory_views():
    """绘制导弹M1三维轨迹 + 三视图投影"""

    # ---------------- 初始参数 ----------------
    x0, y0, z0 = 20000, 0, 2000
    target_x, target_y, target_z = 0, 0, 0

    v_m = 300
    distance = np.sqrt((x0 - target_x) ** 2 + (y0 - target_y) ** 2 + (z0 - target_z) ** 2)
    t_hit = distance / v_m  # ≈ 67.08 s

    key_times = [25, 60, t_hit]
    key_labels = [
        r"$t=25\,\mathrm{s}$：真目标进入视场",
        r"$t=60\,\mathrm{s}$：真目标脱离视场",
        fr"$t\approx{t_hit:.2f}\,\mathrm{{s}}$：击中假目标"
    ]

    t_total = t_hit + 5
    t_values = np.linspace(0, t_total, 300)

    x = x0 - (x0 / t_hit) * t_values
    y = y0 - (y0 / t_hit) * t_values
    z = z0 - (z0 / t_hit) * t_values

    # ---------------- 绘制主图 (3D) ----------------
    fig = plt.figure(figsize=(15, 12))
    ax3d = fig.add_subplot(221, projection="3d")

    ax3d.plot(x, y, z, "k-", linewidth=2.5, label="导弹M1轨迹")

    colors = ["green", "orange", "red"]
    markers = ["o", "s", "D"]

    for i, t in enumerate(key_times):
        t_idx = np.argmin(np.abs(t_values - t))
        ax3d.scatter(x[t_idx], y[t_idx], z[t_idx], color=colors[i],
                     s=120, marker=markers[i], edgecolors="black", linewidth=1.2)
        ax3d.text(x[t_idx] + 200, y[t_idx] + 200, z[t_idx] + 200,
                  key_labels[i], color=colors[i], fontsize=10, weight="bold")

    ax3d.scatter(x0, y0, z0, color="purple", s=150, marker="*",
                 edgecolors="black", linewidth=1.5, label="初始位置")
    ax3d.scatter(target_x, target_y, target_z, color="gold", s=150, marker="X",
                 edgecolors="black", linewidth=1.5, label="假目标")

    ax3d.set_xlabel("$x$ 坐标 (m)")
    ax3d.set_ylabel("$y$ 坐标 (m)")
    ax3d.set_zlabel("$z$ 坐标 (m)")
    ax3d.set_title("导弹M1三维轨迹", fontsize=16)
    ax3d.view_init(elev=22, azim=40)
    ax3d.legend()

    # ---------------- 三视图 ----------------
    # 1. x-y 平面
    ax_xy = fig.add_subplot(222)
    ax_xy.plot(x, y, "k-", linewidth=2)
    for i, t in enumerate(key_times):
        t_idx = np.argmin(np.abs(t_values - t))
        ax_xy.scatter(x[t_idx], y[t_idx], color=colors[i], s=60, marker=markers[i], edgecolors="black")
    ax_xy.scatter(x0, y0, color="purple", s=80, marker="*")
    ax_xy.scatter(target_x, target_y, color="gold", s=80, marker="X")
    ax_xy.set_xlabel("$x$ (m)")
    ax_xy.set_ylabel("$y$ (m)")
    ax_xy.set_title("投影：x-y 平面")
    ax_xy.grid(True)

    # 2. y-z 平面
    ax_yz = fig.add_subplot(223)
    ax_yz.plot(y, z, "k-", linewidth=2)
    for i, t in enumerate(key_times):
        t_idx = np.argmin(np.abs(t_values - t))
        ax_yz.scatter(y[t_idx], z[t_idx], color=colors[i], s=60, marker=markers[i], edgecolors="black")
    ax_yz.scatter(y0, z0, color="purple", s=80, marker="*")
    ax_yz.scatter(target_y, target_z, color="gold", s=80, marker="X")
    ax_yz.set_xlabel("$y$ (m)")
    ax_yz.set_ylabel("$z$ (m)")
    ax_yz.set_title("投影：y-z 平面")
    ax_yz.grid(True)

    # 3. x-z 平面
    ax_xz = fig.add_subplot(224)
    ax_xz.plot(x, z, "k-", linewidth=2)
    for i, t in enumerate(key_times):
        t_idx = np.argmin(np.abs(t_values - t))
        ax_xz.scatter(x[t_idx], z[t_idx], color=colors[i], s=60, marker=markers[i], edgecolors="black")
    ax_xz.scatter(x0, z0, color="purple", s=80, marker="*")
    ax_xz.scatter(target_x, target_z, color="gold", s=80, marker="X")
    ax_xz.set_xlabel("$x$ (m)")
    ax_xz.set_ylabel("$z$ (m)")
    ax_xz.set_title("投影：x-z 平面")
    ax_xz.grid(True)

    # ---------------- 图注 ----------------
    plt.figtext(0.5, 0.01,
                "图注：左上角为三维轨迹，右上角为 x-y 投影，左下角为 y-z 投影，右下角为 x-z 投影，"
                "有助于分别观察导弹在水平、侧视和俯视下的运动特征。",
                ha="center", fontsize=12,
                bbox=dict(facecolor="white", alpha=0.8, boxstyle="round,pad=0.5"))

    plt.tight_layout()
    plt.savefig("missile_m1_trajectory_with_three_views.png", dpi=300, bbox_inches="tight")
    plt.show()


if __name__ == "__main__":
    plot_missile_m1_trajectory_views()
